from transformers import BertTokenizer, BertModel

bert_tokenizer = BertTokenizer.from_pretrained("pretrain_models/bert")
bert_model = BertModel.from_pretrained("pretrain_models/bert")
word2id_bert = bert_tokenizer.get_vocab()
id2word_bert = {word2id_bert[key]: key for key in word2id_bert}


def get_sentence_vector(text):
    encoded_input = bert_tokenizer(text, return_tensors='pt')
    # print(encoded_input)
    # print(bert_model)
    outputs = bert_model(**encoded_input)
    #
    # # print(dir(outputs))
    word_embeddings = outputs.last_hidden_state[:, 0, :]
    # print(outputs.hidden_statess)
    # print(word_embeddings.shape)
    # output = bert_model.get_sequence_output()
    # print(output)
    return word_embeddings
